Support mixed-form sequence lengths in SDPA forward (cuDNN 9.26+)#430
Open
egilliam-nv wants to merge 1 commit into
Open
Support mixed-form sequence lengths in SDPA forward (cuDNN 9.26+)#430egilliam-nv wants to merge 1 commit into
egilliam-nv wants to merge 1 commit into
Conversation
Allow the Q and KV sides of SDPA forward to independently choose their sequence length representation: per-batch (seq_len_*) or cumulative (cu_seq_len_*). Previously the two forms were mutually exclusive across the whole operation, forcing paged-attention integrations -- whose natural state is a cumulative Q prefix sum alongside per-batch KV lengths -- to materialize a KV-side prefix sum purely to satisfy the pairing rule. - sdpa_support_surface: validation is now per-side (exactly one form per side, sides provided together); the paged requirement accepts either form per side; mixed forms are gated on cuDNN 9.26+ in the UNIFIED surface (the initial public 9.25 release predates the backend-side support, and a 9.25 build that has it cannot be told apart from one that does not). Hoist a has_input helper for the presence checks. - diagonal_band_mask: relax the node's metadata seqlen check to per-side exclusivity, matching the SDPA rules. - tests: per-side form control in the SDPA test harness (cu_seq_len_sides) and a deterministic mixed-form test with non-uniform lengths, including bottom-right causal cases that guard the backend's band-mask alignment derivation. Validated on H100 against a cuDNN 9.27 dev build: mixed forms in both directions match the both-per-batch reference for non-paged, ragged, and paged (ragged Q + page tables, cu_seq_len_q + seq_len_kv) graphs; the deterministic tests skip below the version gate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughSDPA validation now supports independently selected explicit or cumulative sequence-length forms for Q and KV. Test configuration, tensor allocation, graph construction, version gating, and parameterized mixed-form coverage are updated accordingly. ChangesMixed sequence-length support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ExecConfig
participant allocate_tensors
participant create_forward_graph
participant SDPA_attributes
ExecConfig->>allocate_tensors: select Q and KV sequence-length forms
allocate_tensors->>create_forward_graph: provide explicit or cumulative tensors
create_forward_graph->>SDPA_attributes: validate mixed-form inputs
SDPA_attributes-->>create_forward_graph: accept or reject configuration
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
egilliam-nv
added
cat-enhancements
mod-frontend
Collaborator
Author
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-430-da67de4 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Allow the Q and KV sides of SDPA forward to independently choose their sequence-length representation: per-batch (
seq_len_*) or cumulative (cu_seq_len_*). Previously the two forms were mutually exclusive across the whole operation.sdpa_support_surface: validation is now per-side (exactly one form per side; both sides provided together); the paged requirement accepts either form per side; mixed forms are gated on cuDNN 9.26+ in the UNIFIED surface. Ahas_inputhelper is hoisted for the presence checks.diagonal_band_mask: the node's metadata seqlen check is relaxed to per-side exclusivity, matching the SDPA rules.cu_seq_len_sides) plus a deterministic mixed-form test (test_sdpa_mixed_seq_len_forms_L0) with non-uniform per-batch lengths, including bottom-right-causal cases that guard the backend's band-mask alignment derivation.Why
Paged-attention integrations naturally hold a cumulative Q prefix sum (
qo_indptr) alongside per-batch KV lengths. The old same-form-both-sides rule forced them to materialize a KV-side prefix sum purely to satisfy the pairing constraint. Permitting mixed forms removes that unnecessary conversion.The 9.26 floor is deliberate: the initial public cuDNN 9.25 release predates the backend-side support, and a later "9.25" build carrying it cannot be distinguished from one without it (
backend_version()reports the same 92500), so mixed forms are only advertised at 9.26+.Related issues
Backend support merged separately (cuDNN backend MRs, internal). Related to #366 (expose
cu_seq_len_q/kvon thesdpa_fp8Python binding).API and compatibility impact
No public API additions or removals. Purely relaxes an existing validation constraint on the
cu_seq_len_*/seq_len_*SDPA inputs; all previously-valid graphs remain valid. New capability (mixed forms) is version-gated to cuDNN 9.26+ and rejected with a clear error below that.Testing
On H100 (sm90) against a cuDNN 9.27 dev build with the backend support:
pytest test/python/test_mhas_v2.py -k mixed_seq_len_forms— 4/4 pass (both mixed directions x {top-left, bottom-right-causal}).pytest test/python/test_sdpa_edge_cases.py— 10/10 pass (no harness regression).pytest test/python/test_mhas_v2.py::test_sdpa_random_fwd_L0(slice) — 203 passed, 9 skipped.cu_seq_len_q+seq_len_kvand bottom-right causal, all match the both-per-batch reference.Below cuDNN 9.26 the mixed-form tests skip via the version gate.
Summary by CodeRabbit
New Features
Bug Fixes
Tests